import sys

from logger import log
from rtp_parser import RtpParser
from g711_parser import G711Parser
from amr_parser import AmrParser, AmrwbParser
from h264_parser import H264Parser
from h265_parser import H265Parser
from  raw_parser import RawParser


class ExtractContext:
    def __init__(self, arg):
        self.args = arg
        self.payloads = None
        if arg.payload:
            self.payloads = arg.payload.split("/")
        self.parserTable = {}
        self.file = None
        if not arg.output:
            log.error("you must assign [output]")
            exit(0)
        else:
            self.file = open(self.args.output, 'wb')
        self.create_parser()

    def __del__(self):
        self.file.close()

    def handle(self, layer, time):
        payload = bytes.fromhex(layer.payload.raw_value)
        if layer.layer_name not in self.payloads:
            idx = -1
        else:
            idx = self.payloads.index(layer.layer_name)

        for p in self.payloads[idx+1:]:
            parser = self.parserTable[p]
            payload = parser.parse(payload)
            if p == self.payloads[-1]:
                parser.dump(self.file)

    def create_parser(self):
        if not self.payloads:
            return
        for payload in self.payloads:
            if payload == "rtp":
                self.parserTable[payload] = RtpParser()
            elif payload == "pcma" or payload == "pcmu":
                self.parserTable[payload] = G711Parser()
            elif payload == "amr":
                self.parserTable[payload] = AmrParser()
            elif payload == "amr-wb":
                self.parserTable[payload] = AmrwbParser()
            elif payload == "h264":
                self.parserTable[payload] = H264Parser()
            elif payload == "h265":
                self.parserTable[payload] = H265Parser()
            elif payload == "raw":
                self.parserTable[payload] = RawParser()
            else:
                log.error("not support this payload: ", payload)
